SparseSoftmaxCrossEntropyWithLogits

计算稀疏标签下的 Softmax 交叉熵损失及梯度。

与 SoftmaxCrossEntropyWithLogits 不同，该算子的标签（Labels）是类别的索引（整数），而不是 One-hot 向量。这通常能节省内存并加速计算。

算法逻辑如下：

  1. Softmax:

\[p_{i,j} = \frac{e^{x_{i,j}}}{\sum_{k} e^{x_{i,k}}}\]

  1. Cross Entropy Loss: 假设 \(y_i\) 为第 \(i\) 个样本的标签索引：

\[loss_i = - \log(p_{i, y_i})\]

  1. Gradients (当 is_grad=1 时):

\[\frac{\partial loss}{\partial x_{i,j}} = p_{i,j} - \mathbb{1}(j == y_i)\]
输入：

  • input - 输入 Logits 地址。形状为 \([N, C]\)。

  • losses - 输出损失地址。

  • sum_data - 中间计算缓冲区（Workspace）。

  • inner_size - 内部维度大小。对于二维输入 \([N, C]\)，通常为 1。

  • outter_size - 外部维度大小。对于二维输入 \([N, C]\)，通常为 \(N\) (Batch Size)。

  • axis_size - 轴维度大小。对于二维输入 \([N, C]\)，通常为 \(C\) (Num Classes)。

  • labels - 标签索引地址。类型为 int，形状为 \([N]\)。

  • is_grad - 是否计算梯度 (0: 否, 1: 是)。

  • output - 输出梯度地址。如果 is_grad 为 1，则存储计算出的梯度，形状同输入。

  • batch_size - 批大小 (N)。

  • number_of_classes - 类别数量 (C)。

  • partial_losses - 中间计算缓冲区（Workspace），用于归约计算。

  • core_mask - 核掩码（仅适用于共享存储版本）。

输出：

  • losses - 计算得到的交叉熵损失。

  • output - 计算得到的梯度（若启用）。

支持平台：

FT78NE MT7004

备注

  • FT78NE 仅支持 Fp32 数据类型。

  • MT7004 支持 Fp32 和 FP16 数据类型。

  • 输入数据必须是二维矩阵。

  • FP16 模式下，labels 依然保持为 int 类型，其余浮点指针变为 float16*。

共享存储版本:

void fp_sparse_softmax_cross_entropy_with_logits_s(float *input, float *losses, float *sum_data, uint64_t inner_size, uint64_t outter_size, uint64_t axis_size, int *labels, uint64_t is_grad, float *output, uint64_t batch_size, uint64_t number_of_classes, float *partial_losses, int core_mask)
void hp_sparse_softmax_cross_entropy_with_logits_s(float16 *input, float16 *losses, float16 *sum_data, uint64_t inner_size, uint64_t outter_size, uint64_t axis_size, int *labels, uint64_t is_grad, float16 *output, uint64_t batch_size, uint64_t number_of_classes, float16 *partial_losses, int core_mask)

C调用示例（FT78NE - Fp32）：

 1#include <stdio.h>
 2#include <stdint.h>
 3
 4int main(int argc, char* argv[]) {
 5    // 假设所有数据位于 DDR 空间
 6    float* input = (float*)0xC0000000;
 7    float* losses = (float*)0xC1000000;
 8    float* sum_data = (float*)0xC2000000;
 9    int* labels = (int*)0xC3000000;
10    float* output = (float*)0xC4000000;
11    float* partial_losses = (float*)0xC5000000;
12
13    uint64_t batch_size = 64;
14    uint64_t number_of_classes = 1000;
15
16    // 针对二维输入 [N, C] 的常用配置
17    uint64_t outter_size = batch_size;
18    uint64_t axis_size = number_of_classes;
19    uint64_t inner_size = 1;
20
21    uint64_t is_grad = 1; // 计算梯度
22    int core_mask = 0xff;
23
24    fp_sparse_softmax_cross_entropy_with_logits_s(input, losses, sum_data,
25                                                  inner_size, outter_size, axis_size,
26                                                  labels, is_grad, output,
27                                                  batch_size, number_of_classes,
28                                                  partial_losses, core_mask);
29
30    return 0;
31}

私有存储版本:

void fp_sparse_softmax_cross_entropy_with_logits_p(float *input, float *losses, float *sum_data, uint64_t inner_size, uint64_t outter_size, uint64_t axis_size, int *labels, uint64_t is_grad, float *output, uint64_t batch_size, uint64_t number_of_classes, float *partial_losses)
void hp_sparse_softmax_cross_entropy_with_logits_p(float16 *input, float16 *losses, float16 *sum_data, uint64_t inner_size, uint64_t outter_size, uint64_t axis_size, int *labels, uint64_t is_grad, float16 *output, uint64_t batch_size, uint64_t number_of_classes, float16 *partial_losses)

C调用示例（MT7004 - FP16）：

 1#include <stdio.h>
 2#include <stdint.h>
 3
 4int main(int argc, char* argv[]) {
 5    // 假设数据位于 AM 空间
 6    float16* input = (float16*)0x10010000;
 7    float16* losses = (float16*)0x10020000;
 8    float16* sum_data = (float16*)0x10030000;
 9    int* labels = (int*)0x10040000;       // 标签仍为 int
10    float16* output = (float16*)0x10050000;
11    float16* partial_losses = (float16*)0x10060000;
12
13    uint64_t batch_size = 32;
14    uint64_t number_of_classes = 100;
15
16    uint64_t outter_size = batch_size;
17    uint64_t axis_size = number_of_classes;
18    uint64_t inner_size = 1;
19
20    uint64_t is_grad = 0; // 仅计算 Loss，不计算梯度
21
22    hp_sparse_softmax_cross_entropy_with_logits_p(input, losses, sum_data,
23                                                   inner_size, outter_size, axis_size,
24                                                   labels, is_grad, output,
25                                                   batch_size, number_of_classes,
26                                                   partial_losses);
27
28    return 0;
29}